import sys
from PyQt5.QtWidgets import QApplication , QMainWindow
from firstMainWin import *
from img2hex1 import *
from hex2img1 import *
from sm4Encode import *
from LSB import *
# import os

class MyMainWindow(QMainWindow, Ui_MainWindow):    
    global list1
    list1 =[]
    def __init__(self,parent=None):        
        super(MyMainWindow, self).__init__(parent)     
        self.setupUi(self)
        # self.open_file1()
        self.pushButton.clicked.connect(self.open_file1)
        self.pushButton_2.clicked.connect(self.open_file2)
        self.pushButton_3.clicked.connect(self.input_key)
        self.pushButton_5.clicked.connect(self.open_file3)
        self.pushButton_4.clicked.connect(self.jiami)


    def open_file1(self):
        fileName = QtWidgets.QFileDialog.getOpenFileName(self,"选取载体图片","e:/")
        self.listWidget.addItem('您所选择的载体图片为'+fileName[0])
        list1.append(fileName[0])
    def open_file2(self):
        fileName = QtWidgets.QFileDialog.getOpenFileName(self,"选取水印图片","e:/")
        self.listWidget_2.addItem('您所选择的水印图片为'+fileName[0])
        list1.append(fileName[0])
    def input_key(self):
        key = QtWidgets.QInputDialog.getText(self,"请输入","请输入SM4密钥(小密温馨提示是32位十六进制哦~)")
        self.listWidget_3.addItem('已经输入SM4密钥，请您妥善保管自己的密钥')
        list1.append(key[0])
    def open_file3(self):
        fileDirectory = QtWidgets.QFileDialog.getExistingDirectory(self,"选择加密图片存储位置","e:/")
        self.listWidget_4.addItem('您所选择的存储位置为'+fileDirectory)
        list1.append(fileDirectory+'/加密图片.bmp')

    def jiami(self):
        filepath1 = 'e:\\mms\\plaintext.txt'
        filepath2 = 'e:\\mms\\cryptotext.txt'
        imgpath1 = 'e:\\mms\\加密水印图片.bmp'
        img1 = Image.open(list1[1]) #水印图片变成十六进制
        width = img1.size[0]
        height = img1.size[1]
        rgb_img = img1.convert('RGB')
        hexcode = getCode(rgb_img)
        with open(filepath1, 'w') as file_object:
            file_object.write(hexcode)
        key = bytes.fromhex(list1[2]) #水印十六进制明文SM4加密
        sm4 = SM4Cipher(key)
        with open(filepath1) as fp1:
            plain = fp1.read()
        fp2 = open(filepath2,'a+')
        for i in range(0, len(plain)-32, 32):
            fp2.write(sm4.encrypt(bytes.fromhex(plain[i:i+32])).hex())
        fp2.write(sm4.encrypt(bytes.fromhex(plain[len(plain)-32:len(plain)])).hex())
        fp2.close()
        with open(filepath2) as fp:
            a = fp.read()
        showImage(a, width, height)
        img2 = Image.open(list1[0]) #载体图片与水印图片进行LSB
        img3 = Image.open(imgpath1)
        rgb_im2 = img2.convert('RGB')
        rgb_im3 = img3.convert('RGB')
        # 将水印的像素值转为文本
        code = getCode1(rgb_im3)
        # 将水印写入图像
        encry(rgb_im2, code, list1[3])
        self.listWidget_5.addItem('加密成功!已输出到指定文件夹!')
        print('加密成功')
    

if __name__ == '__main__':     
    app = QApplication(sys.argv)
    myWin = MyMainWindow()
    myWin.show()
    sys.exit(app.exec_())
